/*
 * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * The Universal Permissive License (UPL), Version 1.0
 *
 * Subject to the condition set forth below, permission is hereby granted to any
 * person obtaining a copy of this software, associated documentation and/or
 * data (collectively the "Software"), free of charge and under any and all
 * copyright rights in the Software, and any and all patent rights owned or
 * freely licensable by each licensor hereunder covering either (i) the
 * unmodified Software as contributed to or provided by such licensor, or (ii)
 * the Larger Works (as defined below), to deal in both
 *
 * (a) the Software, and
 *
 * (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
 * one is included with the Software each a "Larger Work" to which the Software
 * is contributed by such licensors),
 *
 * without restriction, including without limitation the rights to copy, create
 * derivative works of, display, perform, and distribute the Software and make,
 * use, sell, offer for sale, import, export, have made, and have sold the
 * Software and the Larger Work(s), and to sublicense the foregoing rights on
 * either these or other terms.
 *
 * This license is subject to the following condition:
 *
 * The above copyright notice and either this complete permission notice or at a
 * minimum a reference to the UPL must be included in all copies or substantial
 * portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package org.graalvm.nativebridge;

import org.graalvm.collections.Pair;

import java.lang.annotation.Annotation;
import java.lang.reflect.Array;
import java.lang.reflect.GenericArrayType;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Objects;

/**
 * Configuration used by classes generated by the native bridge processor to look up custom type
 * marshallers.
 *
 * @see GenerateHotSpotToNativeBridge
 * @see GenerateHotSpotToNativeBridge
 * @see GenerateProcessToProcessBridge
 * @since 25.0
 */
public final class MarshallerConfig {

    private final Map<Type, BinaryMarshaller<?>> binaryMarshallers;
    private final Map<Class<? extends Annotation>, List<Pair<Class<?>, BinaryMarshaller<?>>>> annotationBinaryMarshallers;

    private MarshallerConfig(Map<Type, BinaryMarshaller<?>> binaryMarshallers,
                    Map<Class<? extends Annotation>, List<Pair<Class<?>, BinaryMarshaller<?>>>> annotationBinaryMarshallers) {
        this.binaryMarshallers = binaryMarshallers;
        this.annotationBinaryMarshallers = annotationBinaryMarshallers;
    }

    /**
     * Looks up {@link BinaryMarshaller} for the {@code type} and {@code annotationTypes}. The
     * method first tries to find a marshaller registered for the {@code type} and some annotation
     * from {@code annotationTypes}. If no such marshaller exists, it tries to find a marshaller
     * registered just for the {@code type}. If there is no such a marshaller it throws the
     * {@link UnsupportedOperationException}.
     *
     * @param type the parameter or return type.
     * @param annotationTypes parameter or method annotation types.
     * @throws UnsupportedOperationException if there is no registered marshaller for the
     *             {@code type}.
     */
    @SuppressWarnings("unchecked")
    @SafeVarargs
    public final <P> BinaryMarshaller<P> lookupMarshaller(Class<P> type, Class<? extends Annotation>... annotationTypes) {
        BinaryMarshaller<?> res = lookupBinaryMarshallerImpl(type, annotationTypes);
        if (res != null) {
            return (BinaryMarshaller<P>) res;
        } else {
            throw unsupported(type);
        }
    }

    /**
     * Looks up {@link BinaryMarshaller} for the {@code parameterizedType} and
     * {@code annotationTypes}. The method first tries to find a marshaller registered for the
     * {@code parameterizedType} and some annotation from {@code annotationTypes}. If no such
     * marshaller exists, it tries to find a marshaller registered just for the
     * {@code parameterizedType}. If there is no such a marshaller it throws the
     * {@link UnsupportedOperationException}.
     *
     * @param parameterizedType the parameter or return type.
     * @param annotationTypes parameter or method annotation types.
     * @throws UnsupportedOperationException if there is no registered marshaller for the
     *             {@code parameterizedType}.
     */
    @SuppressWarnings("unchecked")
    @SafeVarargs
    public final <P> BinaryMarshaller<P> lookupMarshaller(TypeLiteral<P> parameterizedType, Class<? extends Annotation>... annotationTypes) {
        BinaryMarshaller<?> res = lookupBinaryMarshallerImpl(parameterizedType.getType(), annotationTypes);
        if (res != null) {
            return (BinaryMarshaller<P>) res;
        } else {
            throw unsupported(parameterizedType.getType());
        }
    }

    /**
     * Creates a new {@link MarshallerConfig} builder.
     */
    public static Builder newBuilder() {
        return new Builder();
    }

    /**
     * Returns a {@link BinaryMarshaller} for stack trace marshalling.
     */
    public static BinaryMarshaller<StackTraceElement[]> defaultStackTraceMarshaller() {
        return DefaultStackTraceMarshaller.INSTANCE;
    }

    public static final class Builder {

        private final Map<Type, BinaryMarshaller<?>> binaryMarshallers;
        private final Map<Class<? extends Annotation>, List<Pair<Class<?>, BinaryMarshaller<?>>>> annotationBinaryMarshallers;

        private Builder() {
            this.binaryMarshallers = new HashMap<>();
            this.annotationBinaryMarshallers = new HashMap<>();
            // Register default marshallers
            this.binaryMarshallers.put(String.class, BinaryMarshaller.nullable(new StringMarshaller()));
            this.binaryMarshallers.put(Throwable.class, new DefaultThrowableMarshaller());
            this.binaryMarshallers.put(StackTraceElement.class, StackTraceElementMarshaller.INSTANCE);
        }

        /**
         * Registers a {@link BinaryMarshaller} for the {@code type}.
         *
         * @param type the type to register {@link BinaryMarshaller} for.
         * @param marshaller the marshaller to register.
         */
        public <P> Builder registerMarshaller(Class<P> type, BinaryMarshaller<P> marshaller) {
            Objects.requireNonNull(type, "Type must be non null.");
            Objects.requireNonNull(marshaller, "Marshaller must be non null.");
            this.binaryMarshallers.put(type, marshaller);
            return this;
        }

        /**
         * Registers a {@link BinaryMarshaller} for the {@code parameterizedType}.
         *
         * @param parameterizedType the type to register {@link BinaryMarshaller} for.
         * @param marshaller the marshaller to register.
         */
        public <P> Builder registerMarshaller(TypeLiteral<P> parameterizedType, BinaryMarshaller<P> marshaller) {
            Objects.requireNonNull(parameterizedType, "ParameterizedType must be non null.");
            Objects.requireNonNull(marshaller, "Marshaller must be non null.");
            this.binaryMarshallers.put(parameterizedType.getType(), marshaller);
            return this;
        }

        /**
         * Registers a {@link BinaryMarshaller} for the {@code type} and {@code annotationType}.
         *
         * @param type the type to register {@link BinaryMarshaller} for.
         * @param annotationType a required annotation to look up the marshaller.
         * @param marshaller the marshaller to register.
         *
         */
        public <P> Builder registerMarshaller(Class<P> type, Class<? extends Annotation> annotationType, BinaryMarshaller<P> marshaller) {
            Objects.requireNonNull(type, "Type must be non null.");
            Objects.requireNonNull(annotationType, "AnnotationType must be non null.");
            Objects.requireNonNull(marshaller, "Marshaller must be non null.");
            insert(annotationBinaryMarshallers, type, annotationType, marshaller);
            return this;
        }

        /**
         * Registers a {@link BinaryMarshaller} for the {@code parameterizedType} and
         * {@code annotationType}.
         *
         * @param parameterizedType the type to register {@link BinaryMarshaller} for.
         * @param annotationType a required annotation to look up the marshaller.
         * @param marshaller the marshaller to register.
         *
         */
        public <P> Builder registerMarshaller(TypeLiteral<P> parameterizedType, Class<? extends Annotation> annotationType, BinaryMarshaller<P> marshaller) {
            Objects.requireNonNull(parameterizedType, "ParameterizedType must be non null.");
            Objects.requireNonNull(annotationType, "AnnotationType must be non null.");
            Objects.requireNonNull(marshaller, "Marshaller must be non null.");
            insert(annotationBinaryMarshallers, parameterizedType.getRawType(), annotationType, marshaller);
            return this;
        }

        /**
         * Builds {@link MarshallerConfig}.
         */
        public MarshallerConfig build() {
            return new MarshallerConfig(binaryMarshallers, annotationBinaryMarshallers);
        }

        private static <T> void insert(Map<Class<? extends Annotation>, List<Pair<Class<?>, T>>> into, Class<?> type, Class<? extends Annotation> annotationType, T marshaller) {
            verifyAnnotation(annotationType);
            List<Pair<Class<?>, T>> types = into.computeIfAbsent(annotationType, (k) -> new LinkedList<>());
            Pair<Class<?>, T> toInsert = Pair.create(type, marshaller);
            boolean inserted = false;
            for (ListIterator<Pair<Class<?>, T>> it = types.listIterator(); it.hasNext();) {
                Pair<Class<?>, T> current = it.next();
                if (current.getLeft().isAssignableFrom(type)) {
                    it.set(toInsert);
                    it.add(current);
                    inserted = true;
                    break;
                }
            }
            if (!inserted) {
                types.add(toInsert);
            }
        }
    }

    @SafeVarargs
    private BinaryMarshaller<?> lookupBinaryMarshallerImpl(Type type, Class<? extends Annotation>... annotationTypes) {
        for (Class<? extends Annotation> annotationType : annotationTypes) {
            verifyAnnotation(annotationType);
            BinaryMarshaller<?> res = lookup(annotationBinaryMarshallers, type, annotationType);
            if (res != null) {
                return res;
            }
        }
        return binaryMarshallers.get(type);
    }

    private static void verifyAnnotation(Class<? extends Annotation> annotationType) {
        if (annotationType.getAnnotation(MarshallerAnnotation.class) == null) {
            throw new IllegalArgumentException(String.format("The %s in not a valid marshaller annotation. The marshaller annotation must be annotated by the %s meta-annotation.",
                            annotationType.getSimpleName(), MarshallerAnnotation.class.getSimpleName()));
        }
    }

    private static <T> T lookup(Map<Class<? extends Annotation>, List<Pair<Class<?>, T>>> marshallers, Type type, Class<? extends Annotation> annotationType) {
        List<Pair<Class<?>, T>> marshallersForAnnotation = marshallers.get(annotationType);
        if (marshallersForAnnotation != null) {
            Class<?> rawType = erasure(type);
            for (Pair<Class<?>, T> marshaller : marshallersForAnnotation) {
                if (marshaller.getLeft().isAssignableFrom(rawType)) {
                    return marshaller.getRight();
                }
            }
        }
        return null;
    }

    private static Class<?> erasure(Type type) {
        if (type instanceof Class) {
            return (Class<?>) type;
        } else if (type instanceof ParameterizedType) {
            return (Class<?>) ((ParameterizedType) type).getRawType();
        } else if (type instanceof GenericArrayType) {
            return arrayTypeFromComponentType(erasure(((GenericArrayType) type).getGenericComponentType()));
        } else {
            throw new IllegalArgumentException("Unsupported type: " + type);
        }
    }

    private static Class<?> arrayTypeFromComponentType(Class<?> componentType) {
        return Array.newInstance(componentType, 0).getClass();
    }

    private static RuntimeException unsupported(Type type) {
        throw new UnsupportedOperationException(String.format("Marshalling of %s is not supported", type));
    }
}
